#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <openssl/evp.h>
#include <openssl/pem.h>

int sm2_sign(const char *privkey_file, const unsigned char *data, size_t data_len,
             unsigned char **sig, size_t *sig_len) {
    FILE *fp = fopen(privkey_file, "r");
    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);
    if (!pkey) return -1;

    EVP_MD_CTX *md_ctx = EVP_MD_CTX_new();
    if (!md_ctx) { EVP_PKEY_free(pkey); return -1; }

    if (EVP_DigestSignInit(md_ctx, NULL, EVP_sm3(), NULL, pkey) <= 0 ||
        EVP_PKEY_CTX_set_ec_scheme(EVP_MD_CTX_pkey_ctx(md_ctx), NID_sm2) <= 0 ||
        EVP_DigestSign(md_ctx, NULL, sig_len, data, data_len) <= 0) {
        EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey); return -1;
    }

    *sig = malloc(*sig_len);
    if (EVP_DigestSign(md_ctx, *sig, sig_len, data, data_len) <= 0) {
        free(*sig); *sig = NULL;
        EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey); return -1;
    }

    EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey);
    return 0;
}

int main() {
    const char *msg = "fsl";
    const char *id = "1234567812345678"; // 16 bytes
    unsigned char *sig = NULL;
    size_t sig_len = 0;

    // Prepend ID as per SM2 standard (Z = SM3(ID || ...))
    // OpenSSL 3.0+ handles ID via EVP_PKEY_CTX_set1_id()
    FILE *fp = fopen("sm2private_key.pem", "r");
    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);
    if (!pkey) { fprintf(stderr, "Load private key failed\n"); return 1; }

    EVP_MD_CTX *md_ctx = EVP_MD_CTX_new();
    EVP_PKEY_CTX *pctx;
    if (EVP_DigestSignInit(md_ctx, &pctx, EVP_sm3(), NULL, pkey) <= 0 ||
        EVP_PKEY_CTX_set_ec_scheme(pctx, NID_sm2) <= 0 ||
        EVP_PKEY_CTX_set1_id(pctx, (unsigned char*)id, strlen(id)) <= 0 ||
        EVP_DigestSignUpdate(md_ctx, msg, strlen(msg)) != 1 ||
        EVP_DigestSignFinal(md_ctx, NULL, &sig_len) <= 0) {
        EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey);
        fprintf(stderr, "Sign init failed\n");
        return 1;
    }

    sig = malloc(sig_len);
    if (EVP_DigestSignFinal(md_ctx, sig, &sig_len) <= 0) {
        free(sig);
        EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey);
        fprintf(stderr, "Sign final failed\n");
        return 1;
    }

    EVP_MD_CTX_free(md_ctx); EVP_PKEY_free(pkey);

    printf("SM2 sign success, signature length: %zu bytes\n", sig_len);
    // 可选：保存签名到文件
    FILE *sig_fp = fopen("sm2.sig", "wb");
    fwrite(sig, 1, sig_len, sig_fp);
    fclose(sig_fp);
    free(sig);
    return 0;
}
